import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
import xgboost as xgb


# 特征工程，给字段值赋予权重
def data_processing(pth):
    data = pd.read_csv(r'..\data\train.csv')
    data.drop(['Over18', 'StandardHours', 'EmployeeNumber', 'Gender'],
              axis='columns', inplace=True)
    data[['Education', 'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel', 'JobSatisfaction',
          'RelationshipSatisfaction',
          'StockOptionLevel', 'WorkLifeBalance']] = data[
        ['Education', 'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel', 'JobSatisfaction',
         'RelationshipSatisfaction', 'StockOptionLevel', 'WorkLifeBalance']].astype(str)
    return data


# 进行模型训练
def train_model(data):
    # %%
    df = pd.get_dummies(data)
    x = df.drop(['Attrition', 'OverTime_No'], axis='columns')
    y = df['Attrition']
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=24)
    es = xgb.XGBClassifier(max_depth=3, random_state=25)
    es.fit(x_train, y_train)
    y_pre = es.predict(x_test)
    y_pre_proba = es.predict_proba(x_test)[:, 1]
    print('roc评分：{}'.format(roc_auc_score(y_test, y_pre_proba)))
    return es


if __name__ == '__main__':
    df = data_processing('../data/train.csv')
    es = train_model(df)
